#
# 前置条件在../shell脚本/获取系统信息.sh
#

import psutil
import time
from datetime import datetime
import subprocess

def get_system_info():
    # 获取 CPU 使用率
    cpu_usage = psutil.cpu_percent(interval=1)

    # 获取内存使用率
    mem_usage = psutil.virtual_memory().percent

    # 获取所有挂载点的磁盘使用率
    disk_usage = []
    for partition in psutil.disk_partitions():
        try:
            usage = psutil.disk_usage(partition.mountpoint)
            disk_usage.append(( partition.mountpoint, usage.percent))
        except PermissionError:
            continue

    # 获取网络速率（最近 5 秒的平均速率）
    net_io_start = psutil.net_io_counters()
    time.sleep(5)  # 等待 5 秒
    net_io_end = psutil.net_io_counters()

    sent_rate = (net_io_end.bytes_sent - net_io_start.bytes_sent) / 5 / 1024 / 1024  # MB/s
    recv_rate = (net_io_end.bytes_recv - net_io_start.bytes_recv) / 5 / 1024 / 1024  # MB/s

    # 获取系统负载
    load_avg = psutil.getloadavg()

    # 获取系统运行时间
    boot_time = datetime.fromtimestamp(psutil.boot_time())
    uptime = datetime.now() - boot_time
    up_seconds = uptime.total_seconds() 
    up_days = round(up_seconds / 86400, 2)

    # 获取温度信息
    try:
        with open('/sys/class/thermal/thermal_zone0/temp', 'r') as f:
            cpu_temp = int(f.read()) / 1000  # 转换为摄氏度
    except FileNotFoundError:
        cpu_temp = None

    # 获取前三 CPU 使用率的进程
    top_cpu_processes = get_top_processes(sort_by="%cpu", limit=3)

    # 获取前三内存使用率的进程
    top_mem_processes = get_top_processes(sort_by="%mem", limit=3)

    # 打印系统信息
    print("===== 系统信息 =====" + "\\n")
    print(f"CPU : {cpu_usage}%" + "\\n")
    print(f"内存: {mem_usage}%" + "\\n")
    for mountpoint, usage in disk_usage:
        print(f"({mountpoint}) : {usage}%" + "\\n")
    print(f"网络发送速率: {sent_rate:.2f} MB/s" + "\\n")
    print(f"网络接收速率: {recv_rate:.2f} MB/s" + "\\n")
    print(f"(load_5): {load_avg[1]}" + "\\n")
    print(f"系统运行时间: {up_days} 天" + "\\n")
    if cpu_temp:
        print(f"CPU 温度: {cpu_temp}°C" + "\\n")
    else:
        print("CPU 温度: 不支持或未检测到" + "\\n")

    # 打印前三 CPU 使用率的进程
    print("\\n前三 CPU 使用率的进程" + "\\n")
    print("CPU% MEM% COMMAND" + "\\n")
    for line in top_cpu_processes:
        print(f"{line}" + "\\n")

    # 打印前三内存使用率的进程
    print("\\n前三内存使用率的进程" + "\\n")
    print("CPU% MEM% COMMAND" + "\\n")
    for line in top_mem_processes:
        print(f"{line}" + "\\n")

def get_top_processes(sort_by, limit=3):
    # 调用 ps 命令获取进程信息
    ps_output = subprocess.check_output(
        ["ps", "-aux", f"--sort=-{sort_by}"],
        universal_newlines=True
    )

    # 按行分割输出
    lines = ps_output.splitlines()

    # 提取前 limit + 1 行（第一行是表头）
    top_processes = lines[:limit + 1]

    # 解析并提取必要信息
    result = []
    for line in top_processes[1:]:  # 跳过表头
        fields = line.split(maxsplit=10)  # 最多分割 10 次，确保命令部分完整
        cpu = fields[2]
        mem = fields[3]
        command = fields[10].split()[0]  # 截取命令到第一个空格
        result.append(f"{cpu}% {mem}% {command}")

    return result

if __name__ == "__main__":
    get_system_info()